import { CvNodeData, ExternalData } from "@/components/workflow/cv-node";
import { config } from "@/config";
import { Node, useReactFlow } from "@xyflow/react";
import { produce } from "immer";

function generateRequestBody(data: CvNodeData): Record<string, unknown> {
  if (data.inputs.length === 0) {
    throw new Error("No inputs provided");
  }
  if (data.inputs.some((item) => item.value === "")) {
    throw new Error("Some inputs are empty");
  }
  if (["preprocessing", "classification", "detection"].includes(data.task)) {
    return {
      task_type: data.task,
      model_name: data.modelName,
      image_base64: data.inputs[0].value,
      label: "-1",
    };
  } else if (data.task === "search") {
    const image_base64 = data.inputs
      .filter((item) => item.type === "image")
      .map((item) => item.value);
    const image_name = data.inputs
      .filter((item) => item.type === "image")
      .map((item) => item.name + ".jpg");
    return {
      task_type: data.task,
      model_name: data.modelName,
      image_base64,
      image_name,
      gallery_txt_content: data.inputs.find(
        (item) => item.type === "text" && item.name === "gallery-content"
      )?.value,
      query_txt_content: data.inputs.find(
        (item) => item.type === "text" && item.name === "query_content"
      )?.value,
    };
  } else if (data.task === "tracking") {
    const image_base64 = data.inputs
      .filter((item) => item.type === "image")
      .map((item) => item.value);
    const init_bbox = data.inputs.find((item) => item.type === "text")?.value;
    return {
      task_type: data.task,
      model_name: data.modelName,
      image_base64,
      init_bbox,
    };
  }
  return {};
}

interface ResponseBody extends Record<string, unknown> {
  result_text: string;
  output_image_base64: [
    {
      filename: string;
      image_base64: string;
    }
  ];
}

function parseResponseBody(
  data: CvNodeData,
  responseBody: ResponseBody
): ExternalData[] {
  if (["preprocessing", "classification", "detection"].includes(data.task)) {
    return produce(data.outputs, (draft) => {
      draft.forEach((output) => {
        if (output.name === "result_information") {
          output.value = responseBody["result_text"];
        } else if (output.name === "result") {
          output.value = responseBody["output_image_base64"][0].image_base64;
        }
      });
    });
  } else if (data.task === "search") {
    const result: ExternalData[] = [];
    result.push(
      ...responseBody["output_image_base64"].map((item) => ({
        name: item.filename,
        type: "image" as "image" | "text",
        value: item.image_base64,
        reference: "",
      }))
    );
    result.push({
      name: "result_information",
      type: "text",
      value: responseBody["result_text"],
      reference: "",
    });
    return result;
  } else if (data.task === "tracking") {
    const result: ExternalData[] = [];
    result.push(
      ...responseBody["output_image_base64"].map((item) => ({
        name: "result",
        type: "image" as "image" | "text",
        value: item.image_base64,
        reference: "",
      }))
    );
    return result;
  }
  return [];
}

export const useNodeExecuter = () => {
  const { getNode, updateNodeData, getEdges } =
    useReactFlow<Node<CvNodeData>>();

  const executeNode = async (
    nodeId: string,
    options?: {
      payload?: ExternalData[];
      recursive?: boolean;
    }
  ) => {
    const { payload, recursive = false } = options || {};
    try {
      const node = getNode(nodeId);
      if (!node) {
        throw new Error(`Node with id ${nodeId} not found`);
      }
      updateNodeData(nodeId, {
        status: "running",
      });

      let requestBody;
      if (payload) {
        const newData = produce(node.data, (draft) => {
          draft.inputs.forEach((input) => {
            if (input.reference) {
              const found = payload.find(
                (item) =>
                  item.name === input.reference &&
                  item.type === input.type &&
                  item.value !== ""
              );
              if (found) {
                input.value = found.value;
              }
            }
          });
        });
        updateNodeData(nodeId, { inputs: newData.inputs });
        requestBody = generateRequestBody(newData);
      } else {
        requestBody = generateRequestBody(node.data);
      }

      const response = await fetch(
        `${config.AI_SERVICE_URL}${config.AI_SERVICE_URL_PREFIX}/run_cv_inference`,
        {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
          },
          body: JSON.stringify(requestBody),
        }
      );
      if (!response.ok) {
        throw new Error(`API call failed with status ${response.status}`);
      }
      const result = await response.json();
      const outputs = parseResponseBody(node.data, result) || [];
      updateNodeData(nodeId, {
        status: "success",
        outputs,
      });

      if (recursive) {
        const edge = getEdges().find((edge) => edge.source === nodeId);
        if (edge) {
          executeNode(edge.target, {
            payload: outputs,
            recursive,
          });
        }
      }
    } catch (error) {
      console.error(`Error executing node ${nodeId}:`, error);
      updateNodeData(nodeId, {
        status: "failed",
      });
    }
  };

  return { executeNode };
};
